import matplotlib.pyplot as plt
from keras.datasets.cifar10 import load_data
from keras import Sequential, layers, activations, optimizers, losses, Model

(x_train, y_train), (x_test, y_test) = load_data()

x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

onehot_dim = len(set(y_train.flatten()))

class Lenet5(Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.conv = Sequential([
            layers.Conv2D(filters=6, kernel_size=(5, 5), activation=activations.relu),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=16, kernel_size=(5, 5), activation=activations.relu),
            layers.MaxPooling2D(),
        ])
        self.flat = layers.Flatten()
        self.fc = Sequential([
            layers.Dense(units=120, activation=activations.relu),
            layers.Dropout(0.3),
            layers.Dense(units=84, activation=activations.relu),
            layers.Dropout(0.3),
            layers.Dense(units=onehot_dim, activation=activations.softmax)
        ])
    def call(self, inputs, training=None, mask=None):
        out = self.conv(inputs)
        out = self.flat(out)
        out = self.fc(out)
        return out

model = Lenet5()
model.build(input_shape=(None, 32, 32, 3))
model.summary()

model.compile(optimizer=optimizers.Adam(), loss=losses.sparse_categorical_crossentropy, metrics='acc')
log = model.fit(x_train, y_train, epochs=10, batch_size=100, validation_data=(x_test, y_test))

train_acc = log.history['acc']
val_acc = log.history['val_acc']

plt.plot(train_acc, color='r')
plt.plot(val_acc, color='g')
plt.show()
